﻿using System;
using System.Collections.Concurrent;
using System.IO;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace ConsoleAppPressureTest
{
    class Program
    {
        static void Main(string[] args)
        {
            BEGIN:
            Output("Input the url of the socket server:");
            var url = Console.ReadLine();
            if (url.IsNullOrEmpty())
            {
                goto BEGIN;
            }
            Output("Input the connections count for test(default by 1000):");
            var strNum = Console.ReadLine();
            var num = strNum.ToInt();
            if (num == 0)
            {
                num = 100;
            }

            Test(url, num).GetAwaiter();

        }
        private static ConcurrentDictionary<WebSocket, int> clients = new ConcurrentDictionary<WebSocket, int>();
        private static CancellationToken token = new CancellationToken();

        static async Task Test(string url, int num)
        {
            for (var i = 0; i < num; ++i)
            {
                ClientWebSocket webSocket = new ClientWebSocket();
                webSocket.Options.UseDefaultCredentials = true;
                webSocket.Options.Cookies = new System.Net.CookieContainer();
                try
                {
                    await webSocket.ConnectAsync(new Uri(url), CancellationToken.None);
                    clients.TryAdd(webSocket, i);
                }catch(Exception e)
                {
                    
                }
            }
            await Task.Run(async () =>
            {
                foreach (var key in clients.Keys)
                {
                    await SendAsync(key, $"[{DateTime.Now.ToString("yyyy-MM-dd HH:mm:ss")}]from client {clients[key]} message");
                    await ReceiveAsync(key, clients[key]);
                }
            });
        }

        static async Task SendAsync(WebSocket socket, string message)
        {
            if (socket.State != WebSocketState.Open)
                return;
            var buffer = Encoding.UTF8.GetBytes(message);
            await socket.SendAsync(buffer: new ArraySegment<byte>(array: buffer, offset: 0, count: buffer.Length),
                    messageType: WebSocketMessageType.Text,
                    endOfMessage: true,
                    cancellationToken: CancellationToken.None).ConfigureAwait(false);
        }

        static async Task<string> ReceiveAsync(WebSocket socket, int index)
        {
            string message = string.Empty;
            if (socket.State == WebSocketState.Open)
            {
                ArraySegment<byte> buffer = new ArraySegment<byte>(new byte[1024]);
                try
                {
                    WebSocketReceiveResult result;
                    using (var ms = new MemoryStream())
                    {
                        do
                        {
                            result = await socket.ReceiveAsync(buffer, CancellationToken.None).ConfigureAwait(false);
                            ms.Write(buffer.Array, buffer.Offset, result.Count);
                        }
                        while (!result.EndOfMessage);
                        ms.Seek(0, SeekOrigin.Begin);
                        using (var reader = new StreamReader(ms, Encoding.UTF8))
                        {
                            message = await reader.ReadToEndAsync().ConfigureAwait(false);
                            Log(message, index);
                        }
                    }
                }
                catch (WebSocketException e)
                {
                    if (e.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
                    {
                        clients.TryRemove(socket, out _);
                        socket.Abort();
                    }
                }
            }
            return message;
        }

        static string path = "logs/";

        static void Log(string message, int index)
        {
            if (!Directory.Exists(path))
            {
                Directory.CreateDirectory(path);
            }
            var filePath = $"{path}test_{index}.log";
            File.AppendAllText(filePath, message);
        }

        static void Output(string msg, bool bNeedRet = false)
        {
            if (bNeedRet)
                Console.WriteLine(msg);
            else
                Console.Write(msg);
        }
    }
}
